#!/usr/bin/env python
# -*- coding: utf-8 -*-


from tornado import gen, stack_context
import tornado.tcpclient

import socket
import functools

class TCPClient(tornado.tcpclient.TCPClient):
    def __init__(self, resolver=None, io_loop=None, connect_timeout=1):
        super(TCPClient, self).__init__(resolver, io_loop)
        self.connect_timeout = connect_timeout

    @gen.coroutine
    def connect(self, host, port, af=socket.AF_UNSPEC, ssl_options=None,
            max_buffer_size=None):
        @gen.coroutine
        def _connect(host, port, af=socket.AF_UNSPEC, ssl_options=None,
                    max_buffer_size=None):
            """Connect to the given host and port.

            Asynchronously returns an `.IOStream` (or `.SSLIOStream` if
            ``ssl_options`` is not None).
            """
            addrinfo = yield self.resolver.resolve(host, port, af)
            connector = tornado.tcpclient._Connector(
                addrinfo, self.io_loop,
                functools.partial(self._create_stream, max_buffer_size))
            af, addr, stream = yield connector.start()
            # TODO: For better performance we could cache the (af, addr)
            # information here and re-use it on sbusequent connections to
            # the same host. (http://tools.ietf.org/html/rfc6555#section-4.2)
            if ssl_options is not None:
                stream = stream.start_tls(False, ssl_options=ssl_options,
                                                server_hostname=host)
            raise gen.Return(stream)
        future = _connect(host, port, af=socket.AF_INET)
        stream = yield gen.with_timeout(self.io_loop.time() + self.connect_timeout, future)
        raise gen.Return(stream)
